#include "lp_norm_v2_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 32;       // 32B对齐
const uint32_t BUFFER_NUM = 2;        // 双缓冲
const uint32_t SLOT_STRIDE = 64 /sizeof(float) ;    

static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    LpNormV2CustomTilingData tiling;
    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    auto coreNum = ascendcPlatform.GetCoreNum();

    auto axis = tiling.get_axis();
    uint32_t rownum = context->GetInputShape(0)->GetStorageShape().GetDim(0);
    uint32_t colnum = context->GetInputShape(0)->GetStorageShape().GetDim(1);
    auto dt = context->GetInputDesc(0)->GetDataType();


    // Based on the input length and the number of inputs, the number of bytes of the input data type is obtained
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t typeLength = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
    uint32_t inputLength = inputNum * typeLength;
    uint32_t inputBytes = inputLength / inputNum;
   
     uint32_t ubDataNumber = 12;
    // The number of 32B data blocks that can be used for each data. DOUBLE BUFFER is already counted here
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;


    // Input data for 32B alignment
    uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
    // There is at least 32B of data on each core, satisfying several settings for several cores. The maximum number of audits is the actual number of audits
    coreNum = (coreNum <  inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
    coreNum = (coreNum >= 1) ? coreNum : 1;
    uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE / coreNum;
    uint32_t tailBlockNum = (inputLengthAlgin32 / BLOCK_SIZE) % coreNum;
    
    // Small chunks are calculated and sliced several times using the number of data on each core
    uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
    // Tail block calculation for small chunks of data
    uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
    smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
    
    // The total length of a large block of data is 32B larger than that of a small block of data
    everyCoreInputBlockNum += 1;
    uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
    uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
    bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;

    // 8. 设置tiling参数
    tiling.set_smallCoreDataNum(smallCoreDataNum);
    tiling.set_bigCoreDataNum(bigCoreDataNum);
    tiling.set_finalBigTileNum(finalBigTileNum);
    tiling.set_finalSmallTileNum(finalSmallTileNum);
    tiling.set_tileDataNum(tileDataNum);
    tiling.set_smallTailDataNum(smallTailDataNum);
    tiling.set_bigTailDataNum(bigTailDataNum);
    tiling.set_tailBlockNum(tailBlockNum);

    std::cout << "numAttrs=" << context->GetAttrs()->GetAttrNum() << std::endl;



    // 9. 传递算子参数（从context获取）
    tiling.set_p(context->GetAttrs()->GetFloat(0)[0]);
    tiling.set_axis(context->GetAttrs()->GetInt(1)[0]);
    tiling.set_epsilon(context->GetAttrs()->GetFloat(2)[0]);
    tiling.set_dataTypeId(dt);
    tiling.set_cols(colnum);
    tiling.set_rows(rownum);



    // 10. 保存tiling数据并设置工作区大小
    context->SetBlockDim(coreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    
    // workspace 大小设置（系统 + 用户）
    uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); // 系统workspace

    // 用户workspace大小（64B 对齐，每个范数一个slot）

    uint32_t workspaceCount = 0;
    if (axis == 0) {
        workspaceCount = 1;                 // 全局规约
    } else if (axis == 1) {
        workspaceCount = tiling.get_cols();     // 按列规约
    } else if (axis == 2) {
        workspaceCount = tiling.get_rows();     // 按行规约
    }
    uint32_t usrSize = workspaceCount * SLOT_STRIDE * sizeof(float); // 用户部分

    // 返回总大小
    size_t* currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = usrSize + sysWorkspaceSize;

    tiling.set_finalBigTileNum(finalBigTileNum);
    tiling.set_finalSmallTileNum(finalSmallTileNum);

    return ge::GRAPH_SUCCESS;
}
}

namespace ge {
// 形状推断：输出形状与输入一致
static ge::graphStatus InferShape(gert::InferShapeContext* context) {
    const gert::Shape* x_shape = context->GetInputShape(0);
    gert::Shape* z_shape = context->GetOutputShape(0);
    *z_shape = *x_shape;
    return GRAPH_SUCCESS;
}

// 数据类型推断：输出类型与输入一致
static ge::graphStatus InferDataType(gert::InferDataTypeContext* context) {
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
}

namespace ops {
class LpNormV2Custom : public OpDef {
public:
    explicit LpNormV2Custom(const char* name) : OpDef(name) {
        // 输入定义
        this->Input("input")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        // 输出定义
        this->Output("output")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        // 算子属性定义

        this->Attr("p")
            .AttrType(REQUIRED)
            .Float(2.0f);

        this->Attr("axis")
            .AttrType(OPTIONAL)
            .Int(0);

        this->Attr("epsilon")
            .AttrType(OPTIONAL)
            .Float(1e-12f);

        this->Attr("data_type_id")
            .AttrType(REQUIRED)
            .Int(0);


        // 关联形状和数据类型推断函数
        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        // 绑定AICore计算配置
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend310b")
            .AddConfig("ascend910b");
    }
};

// 注册算子
OP_ADD(LpNormV2Custom);
}